import os
import pandas as pd
import logging
from typing import Optional, Tuple, Dict, Any

def setup_logging(log_level: str = 'INFO') -> None:
    """设置日志系统
    
    参数:
        log_level: 日志级别，如 'INFO', 'DEBUG' 等
    """
    # 创建日志记录器
    logger = logging.getLogger('country_attribute_refixer')
    logger.setLevel(log_level.upper())
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(log_level.upper())
    console_formatter = logging.Formatter('%(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 添加处理器
    logger.addHandler(console_handler)
    
    # 设置全局日志记录器
    logging.getLogger().addHandler(console_handler)

def attach_country_attributes(
    result_df: pd.DataFrame, 
    output_path: str, 
    data_dir: str, 
    force_merge: bool = False
) -> Optional[pd.DataFrame]:
    """
    给结果文件附加国家属性数据，并保存新的 CSV 文件

    参数:
        result_df (pd.DataFrame): 排序结果 DataFrame（需包含 '国名Ch' 列）
        output_path (str): 原始结果 CSV 文件路径
        data_dir (str): 国家属性文件所在目录，应包含 'countryAttr.csv'
        force_merge (bool): 强制合并即使用户选择否，默认 False
    
    返回:
        enriched_df (pd.DataFrame or None): 成功附加后返回新的 DataFrame，失败返回 None
    """
    # 设置日志记录器
    setup_logging()
    
    # 验证输入数据
    if '国名Ch' not in result_df.columns:
        logging.error("错误：结果数据框中缺少 '国名Ch' 列，无法进行国家属性附加")
        return None
    
    attr_file = os.path.join(data_dir, "countryAttr.csv")
    required_cols = ['国名Ch', '地理区域', '收入水平', '宗教信仰']
    
    # 检查并处理国家属性文件
    if not os.path.exists(attr_file):
        logging.warning(f"警告：未找到国家属性文件 '{attr_file}'，附加操作已跳过")
        return None
    
    # 询问用户是否附加属性
    # 如果 force_merge 为 True，自动进行合并
    if not force_merge:
        enrich = input("\n是否附加国家属性（地理区域、收入水平、宗教信仰等）？[Y/N，默认Y]: ").strip().lower()
        if enrich not in ['y', 'yes', ''] and not force_merge:
            logging.info("用户已选择跳过附加国家属性步骤")
            return None
    
    try:
        # 读取属性文件
        attr_df = pd.read_csv(attr_file)
        
        # 验证属性文件所需列的存在
        missing_cols = [col for col in required_cols if col not in attr_df.columns]
        if missing_cols:
            logging.error(f"错误：属性文件 '{attr_file}' 缺少必要列: {missing_cols}")
            return None
        
        # 检查国家名称匹配
        result_countries = set(result_df['国名Ch'].astype(str))
        attr_countries = set(attr_df['国名Ch'].astype(str))
        
        missing_in_attr = result_countries - attr_countries
        if missing_in_attr:
            logging.warning(f"警告：以下国家在属性文件中找不到匹配: {missing_in_attr}")
        
        # 创建完整数据
        enriched_df = result_df.merge(attr_df, on='国名Ch', how='left', suffixes=('', '_attr'))
        
        # 处理合并后的问题列
        problematic_cols = {
            '收入水平_attr': '收入水平',
            '收入水平': ('收入水平_attr', '收入水平_attr')
        }
        
        for attr_col, target_cols in problematic_cols.items():
            if attr_col in enriched_df.columns:
                if isinstance(target_cols, str):
                    # 移除命名冲突
                    if attr_col in enriched_df.columns:
                        target_col = target_cols
                        if target_col in enriched_df.columns:
                            enriched_df[target_col] = enriched_df[attr_col]
                            logging.info(f"正在修复列名匹配问题: '{attr_col}' -> '{target_col}'")
                        else:
                            logging.error(f"错误：无法修复列名 '{attr_col}'")
                    enriched_df.drop(columns=[attr_col], inplace=True)
                else:
                    # 处理多个重名列的情况
                    if all(col in enriched_df.columns for col in target_cols):
                        # 找出最长的列名作为目标列
                        target_cols.sort(key=len, reverse=True)
                        
                        # 如果第一个目标列存在，则使用它
                        if target_cols[0] in enriched_df.columns:
                            value_col = enriched_df[target_cols[0]]
                            for col in target_cols:
                                if col in enriched_df.columns and col != target_cols[0]:
                                    value_col += enriched_df[col].fillna('')
                        else:
                            logging.error(f"错误：无法修复列名 '{target_cols}'")
                    else:
                        logging.error(f"错误：在修复列名时，未找到所有需要的列")
        
        # 保存并返回
        before_rows = result_df.shape[0]
        after_rows = enriched_df.shape[0]
        
        if after_rows < before_rows:
            missing_count = before_rows - after_rows
            matching_rate = (after_rows / before_rows) * 100
        
        enriched_output = output_path.replace(".csv", "_附加属性.csv")
        enriched_df.to_csv(enriched_output, index=False, encoding='utf-8-sig')
        
        if after_rows < before_rows:
            logging.warning(f"警告：成功附加属性后，行数从 {before_rows} 减少到 {after_rows} ({missing_count} 行缺失)")
        
        logging.info(f"国家属性附加完成，已保存至: {enriched_output}")
        return enriched_df
    
    except FileNotFoundError:
        logging.error(f"错误：找不到属性文件 '{attr_file}'")
        return None
    except pd.errors.EmptyDataError:
        logging.error(f"错误：属性文件 '{attr_file}' 为空或格式无效")
        return None
    except pd.errors.ParserError:
        logging.error(f"错误：无法解析属性文件 '{attr_file}'")
        return None
    except Exception as e:
        logging.error(f"错误：附加国家属性时发生未知错误: {str(e)}")
        import traceback
        logging.error(traceback.format_exc())
        return None

